package main

import (
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"

	"gopkg.in/yaml.v2"
	"gorm.io/driver/mysql"
	"gorm.io/gen"
	"gorm.io/gorm"
)

// 配置结构体
type Config struct {
	Database struct {
		Host     string `yaml:"host"`
		Port     int    `yaml:"port"`
		User     string `yaml:"user"`
		Password string `yaml:"password"`
		Database string `yaml:"database"`
		Charset  string `yaml:"charset"`
	} `yaml:"database"`

	Generator struct {
		OutputPath        string `yaml:"output_path"`
		QueryFile         string `yaml:"query_file"`
		ModelPackage      string `yaml:"model_package"`
		WithUnitTest      bool   `yaml:"with_unit_test"`
		FieldNullable     bool   `yaml:"field_nullable"`
		FieldWithIndexTag bool   `yaml:"field_with_index_tag"`
		FieldWithTypeTag  bool   `yaml:"field_with_type_tag"`
	} `yaml:"generator"`

	Tables struct {
		Include     []string          `yaml:"include"`
		Exclude     []string          `yaml:"exclude"`
		NameMapping map[string]string `yaml:"name_mapping"`
	} `yaml:"tables"`

	Fields struct {
		TypeMapping  map[string]string `yaml:"type_mapping"`
		IgnoreFields []string          `yaml:"ignore_fields"`
		Tags         map[string]string `yaml:"tags"`
	} `yaml:"fields"`

	Methods struct {
		GenerateCRUD        bool `yaml:"generate_crud"`
		GenerateQuery       bool `yaml:"generate_query"`
		GenerateAssociation bool `yaml:"generate_association"`
	} `yaml:"methods"`
}

func main() {
	// 加载配置
	config := loadConfig()

	// 连接数据库
	db := connectDatabase(config)

	// 创建生成器
	g := createGenerator(config)
	g.UseDB(db)

	// 生成模型
	generateModels(g, db, config)

	// 执行生成
	g.Execute()

	fmt.Println("🎉 GORM模型生成完成！")
	fmt.Printf("📁 输出目录: %s\n", config.Generator.OutputPath)

	// 生成额外的辅助文件
	generateHelperFiles(config)
}

// 加载配置文件
func loadConfig() *Config {
	configFile := "tools/gen_config.yaml"
	if len(os.Args) > 1 {
		configFile = os.Args[1]
	}

	data, err := ioutil.ReadFile(configFile)
	if err != nil {
		log.Fatalf("❌ 读取配置文件失败: %v", err)
	}

	var config Config
	if err := yaml.Unmarshal(data, &config); err != nil {
		log.Fatalf("❌ 解析配置文件失败: %v", err)
	}

	// 应用环境变量覆盖
	applyEnvOverrides(&config)

	fmt.Printf("✅ 配置加载成功: %s\n", configFile)
	return &config
}

// 应用环境变量覆盖
func applyEnvOverrides(config *Config) {
	if host := os.Getenv("DB_HOST"); host != "" {
		config.Database.Host = host
	}
	if port := os.Getenv("DB_PORT"); port != "" {
		fmt.Sscanf(port, "%d", &config.Database.Port)
	}
	if user := os.Getenv("DB_USER"); user != "" {
		config.Database.User = user
	}
	if password := os.Getenv("DB_PASSWORD"); password != "" {
		config.Database.Password = password
	}
	if database := os.Getenv("DB_NAME"); database != "" {
		config.Database.Database = database
	}
}

// 连接数据库
func connectDatabase(config *Config) *gorm.DB {
	dsn := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=%s&parseTime=True&loc=Local",
		config.Database.User,
		config.Database.Password,
		config.Database.Host,
		config.Database.Port,
		config.Database.Database,
		config.Database.Charset,
	)

	db, err := gorm.Open(mysql.Open(dsn), &gorm.Config{})
	if err != nil {
		log.Fatalf("❌ 数据库连接失败: %v", err)
	}

	fmt.Printf("✅ 数据库连接成功: %s@%s:%d/%s\n",
		config.Database.User,
		config.Database.Host,
		config.Database.Port,
		config.Database.Database,
	)

	return db
}

// 创建生成器
func createGenerator(config *Config) *gen.Generator {
	// 确保输出目录存在
	if err := os.MkdirAll(config.Generator.OutputPath, 0755); err != nil {
		log.Fatalf("❌ 创建输出目录失败: %v", err)
	}

	return gen.NewGenerator(gen.Config{
		OutPath:           config.Generator.OutputPath,
		OutFile:           config.Generator.QueryFile,
		ModelPkgPath:      config.Generator.ModelPackage,
		WithUnitTest:      config.Generator.WithUnitTest,
		FieldNullable:     config.Generator.FieldNullable,
		FieldWithIndexTag: config.Generator.FieldWithIndexTag,
		FieldWithTypeTag:  config.Generator.FieldWithTypeTag,
	})
}

// 生成模型
func generateModels(g *gen.Generator, db *gorm.DB, config *Config) {
	// 获取要生成的表
	tables := getTargetTables(db, config)

	fmt.Printf("📋 将生成 %d 个表的模型:\n", len(tables))
	for i, table := range tables {
		fmt.Printf("  %d. %s\n", i+1, table)
	}

	// 为每个表生成模型
	var models []interface{}
	for _, tableName := range tables {
		fmt.Printf("🔄 生成模型: %s\n", tableName)

		// 创建模型配置
		modelOpts := createModelOptions(tableName, config)
		model := g.GenerateModel(tableName, modelOpts...)

		models = append(models, model)
	}

	// 应用模型
	g.ApplyBasic(models...)
}

// 获取目标表列表
func getTargetTables(db *gorm.DB, config *Config) []string {
	// 获取所有表
	allTables := getAllTables(db)

	var targetTables []string

	// 如果指定了包含列表，只生成指定的表
	if len(config.Tables.Include) > 0 {
		for _, table := range config.Tables.Include {
			if contains(allTables, table) {
				targetTables = append(targetTables, table)
			}
		}
	} else {
		// 否则生成所有表，但排除指定的表
		for _, table := range allTables {
			if !contains(config.Tables.Exclude, table) && !isSystemTable(table) {
				targetTables = append(targetTables, table)
			}
		}
	}

	return targetTables
}

// 获取所有表名
func getAllTables(db *gorm.DB) []string {
	var tables []string

	rows, err := db.Raw("SHOW TABLES").Rows()
	if err != nil {
		log.Fatalf("❌ 获取表列表失败: %v", err)
	}
	defer rows.Close()

	for rows.Next() {
		var tableName string
		if err := rows.Scan(&tableName); err != nil {
			continue
		}
		tables = append(tables, tableName)
	}

	return tables
}

// 判断是否为系统表
func isSystemTable(tableName string) bool {
	systemPrefixes := []string{
		"information_schema",
		"performance_schema",
		"mysql",
		"sys",
	}

	for _, prefix := range systemPrefixes {
		if strings.HasPrefix(tableName, prefix) {
			return true
		}
	}

	return false
}

// 创建模型选项
func createModelOptions(tableName string, config *Config) []gen.ModelOpt {
	var opts []gen.ModelOpt

	// 应用表名映射
	if structName, exists := config.Tables.NameMapping[tableName]; exists {
		opts = append(opts, gen.FieldType("", structName))
	}

	// 忽略指定字段
	for _, field := range config.Fields.IgnoreFields {
		opts = append(opts, gen.FieldIgnore(field))
	}

	// 应用字段类型映射
	for field, fieldType := range config.Fields.TypeMapping {
		opts = append(opts, gen.FieldType(field, fieldType))
	}

	// 应用字段标签
	for field, tag := range config.Fields.Tags {
		opts = append(opts, gen.FieldTag(field, tag))
	}

	return opts
}

// 生成辅助文件
func generateHelperFiles(config *Config) {
	// 生成模型初始化文件
	generateModelInit(config)

	// 生成数据库连接文件
	generateDBConnection(config)
}

// 生成模型初始化文件
func generateModelInit(config *Config) {
	content := `package db_model

import (
	"gorm.io/gorm"
)

// DB 全局数据库实例
var DB *gorm.DB

// SetDB 设置数据库实例
func SetDB(db *gorm.DB) {
	DB = db
}

// GetDB 获取数据库实例
func GetDB() *gorm.DB {
	return DB
}
`

	filePath := filepath.Join(config.Generator.OutputPath, "init.go")
	if err := ioutil.WriteFile(filePath, []byte(content), 0644); err != nil {
		log.Printf("⚠️  生成初始化文件失败: %v", err)
	} else {
		fmt.Printf("✅ 生成初始化文件: %s\n", filePath)
	}
}

// 生成数据库连接文件
func generateDBConnection(config *Config) {
	content := fmt.Sprintf(`package db_model

import (
	"fmt"
	"gorm.io/driver/mysql"
	"gorm.io/gorm"
)

// ConnectDB 连接数据库
func ConnectDB() (*gorm.DB, error) {
	dsn := fmt.Sprintf("%%s:%%s@tcp(%%s:%%d)/%%s?charset=%%s&parseTime=True&loc=Local",
		"%s", // user
		"%s", // password  
		"%s", // host
		%d,   // port
		"%s", // database
		"%s", // charset
	)
	
	return gorm.Open(mysql.Open(dsn), &gorm.Config{})
}
`,
		config.Database.User,
		config.Database.Password,
		config.Database.Host,
		config.Database.Port,
		config.Database.Database,
		config.Database.Charset,
	)

	filePath := filepath.Join(config.Generator.OutputPath, "connection.go")
	if err := ioutil.WriteFile(filePath, []byte(content), 0644); err != nil {
		log.Printf("⚠️  生成连接文件失败: %v", err)
	} else {
		fmt.Printf("✅ 生成连接文件: %s\n", filePath)
	}
}

// 辅助函数：检查切片是否包含元素
func contains(slice []string, item string) bool {
	for _, s := range slice {
		if s == item {
			return true
		}
	}
	return false
}
